package com.huqing.icu.service.impl;

import com.huqing.icu.dao.ChatMessageDao;
import com.huqing.icu.dao.ChatSessionDao;
import com.huqing.icu.enums.BusinessConstants;
import com.huqing.icu.enums.ChatRoleEnum;
import com.huqing.icu.exceptions.BusinessException;
import com.huqing.icu.model.ChatMessage;
import com.huqing.icu.model.ChatSession;
import com.huqing.icu.req.app.chatmessage.ChatCompletionReq;
import com.huqing.icu.service.IChatMessageService;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.metadata.EmptyUsage;
import org.springframework.ai.chat.metadata.Usage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.concurrent.atomic.AtomicReference;

/**
 * @Description 对话消息接口实现类
 * @Author huqing
 * @Date 2025/8/20 9:05
 **/
@Service
public class ChatMessageServiceImpl implements IChatMessageService {

    //private static final Logger logger = LoggerFactory.getLogger(IChatMessageService.class);

    @Autowired
    private ChatClient chatClient;

    @Autowired
    private ChatSessionDao chatSessionDao;
    @Autowired
    private ChatMessageDao chatMessageDao;


    /**
     * AI对话
     *
     * @param chatCompletionReq
     * @param userId
     * @return
     */
    @Override
    public Flux<ChatResponse> completion(ChatCompletionReq chatCompletionReq, Long userId) {


        //1、保存用户发送的消息
        //Long userMessageId = null;
        Long userMessageId = saveChatMessage(userId, chatCompletionReq.getChatSessionId(), null, chatCompletionReq.getPrompt(), null, ChatRoleEnum.ASSISTANT.getCode());

        //token使用量
        AtomicReference<Integer> promptTokens = null;
        AtomicReference<Integer> completionTokens = null;
        StringBuilder assistantMessage = new StringBuilder();
        //前端传给我的会话ID，需要传给advisors，让advisors去管理，也就是添加会话ID到AdvisorsContext上下文中,根据不同的会话ID做用户会话的区分
        return chatClient.prompt().user(chatCompletionReq.getPrompt())
                .advisors(o -> o.param(ChatMemory.CONVERSATION_ID, chatCompletionReq.getChatSessionId()))
                .stream().chatResponse()
                .doOnNext(chatResponse -> {
                    //累加，把流示返回的内容拼接成完成的字符串
                    assistantMessage.append(chatResponse.getResult().getOutput().getText());

                    //获取token使用量
                    Usage usage = chatResponse.getMetadata().getUsage();
                    if (usage == null || EmptyUsage.class.equals(usage.getClass())) {
                        System.out.println("111");
                    } else {

                        promptTokens.set(usage.getPromptTokens());
                        completionTokens.set(usage.getCompletionTokens());
                    }
                })
                /*.doOnNext(content -> assistantMessage.append(content))*/
                //流结束，大模型生成的消息落库
                .doOnComplete(() -> {

                    System.out.println(assistantMessage.toString());
                    //保存大模型消息
                    saveChatMessage(userId, chatCompletionReq.getChatSessionId(), userMessageId, assistantMessage.toString(), completionTokens.get(), ChatRoleEnum.ASSISTANT.getCode());

                    //更新用户消息的token使用量
                    chatMessageDao.updateToken(userMessageId, promptTokens.get());
                });
    }


    /**
     * 保存聊天消息
     *
     * @param userId
     * @param chatSessionId
     * @param parentMessageId
     * @param content
     * @param role
     * @return
     */
    private Long saveChatMessage(Long userId, Long chatSessionId, Long parentMessageId, String content, Integer token, String role) {

        //查询对话是否存在
        ChatSession chatSession = chatSessionDao.selectById(chatSessionId);
        if (chatSession == null) {
            throw new BusinessException(BusinessConstants.DATA_NOT_EXIST.getCode(), "对话不存在，请检查");
        }

        //保存消息
        ChatMessage chatMessage = new ChatMessage();
        chatMessage.setParentId(parentMessageId);
        chatMessage.setUserId(userId);
        chatMessage.setChatSessionId(chatSessionId);
        chatMessage.setContent(content);
        chatMessage.setToken(token);
        chatMessage.setRole(role);
        chatMessageDao.saveOrUpdate(chatMessage);
        return chatMessage.getId();
    }
}
